import matplotlib.pyplot as plt

from core.config import test_mode
from core.cuda import cuda_module, cuda
from core.data.dataloader import DataLoader
from implement.datasource.spiral import Spiral
from implement.layers.basic.two_layers_net import TwoLayersNet
from implement.optimizers.sgd import SGD
from utils.functions_collect import softmax_cross_entropy

# 设置超参数
max_epoch = 100
batch_size = 30
hidden_size = 10
lr = 1.0

# 创建数据集和数据加载器
train_set = Spiral(train=True)
test_set = Spiral(train=False)
train_loader = DataLoader(train_set, batch_size)
test_loader = DataLoader(test_set, batch_size, shuffle=False)

# 创建模型和优化器
model = TwoLayersNet(hidden_size, 3)
optimizer = SGD(lr).setup(model)

# 训练过程
for epoch in range(max_epoch):
    for x, t in train_loader:
        y = model(x)
        loss = softmax_cross_entropy(y, t)
        model.cleargrads()
        loss.backward()
        optimizer.update()
    if epoch % 10 == 0:
        print('loss:', loss.data)

# 可视化
x = cuda_module.array([example[0] for example in train_set])
t = cuda_module.array([example[1] for example in train_set])
h = 0.001
x_min, x_max = x[:, 0].min() - .1, x[:, 0].max() + .1
y_min, y_max = x[:, 1].min() - .1, x[:, 1].max() + .1
xx, yy = cuda_module.meshgrid(cuda_module.arange(x_min, x_max, h), cuda_module.arange(y_min, y_max, h))
X = cuda_module.c_[xx.ravel(), yy.ravel()]

# 使用测试模式获取模型预测结果
with test_mode():
    score = model(X)
predict_cls = cuda_module.argmax(score.data, axis=1)
Z = predict_cls.reshape(xx.shape)
xx = cuda.to_numpy(xx)
yy = cuda.to_numpy(yy)
Z = cuda.to_numpy(Z)
plt.contourf(xx, yy, Z)

# 数据点可视化
N, CLS_NUM = 100, 3
markers = ['o', 'x', '^']
colors = ['orange', 'blue', 'green']
for i in range(len(x)):
    c = int(t[i])
    Z = cuda.to_numpy(Z)
    Z = cuda.to_numpy(Z)
    plt.scatter(cuda.to_numpy(x[i][0]), cuda.to_numpy(x[i][1]), s=40, marker=markers[c], c=colors[c])

# 显示图形
plt.show()
